import ptan
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical


class RNDAgent(ptan.agent.BaseAgent):
    # TODO: unify code with DQNAgent, as only action selector is differs.
    def __init__(self, model, action_selector=ptan.actions.ProbabilityActionSelector(), device="cpu",
                 apply_softmax=False, preprocessor=ptan.agent.default_states_preprocessor):
        self.model = model
        self.action_selector = action_selector
        self.device = device
        self.apply_softmax = apply_softmax
        self.preprocessor = preprocessor

    @torch.no_grad()
    def __call__(self, states, agent_states=None):

        if self.preprocessor is not None:
            states = self.preprocessor(states)
            if torch.is_tensor(states):
                states = states.to(self.device)
        dist, int_value, ext_value, probs = self.model(states)
        if self.apply_softmax:
            probs = F.softmax(probs, dim=1)
        probs = probs.data.cpu().numpy()
        actions = self.action_selector(probs)
        return np.array(actions), [(dist, int_value, ext_value, dist.log_prob(torch.Tensor(actions).to(self.device)))]


class RNDModel(nn.Module):

    def __init__(self, obs_size, act_size, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.conv_encoder = nn.Sequential(
            nn.Conv2d(obs_size[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )

        self.conv_out_size = self._get_conv_out(obs_size)
        self.fc_encoder = nn.Sequential(
            nn.Linear(self.conv_out_size, 256),
            nn.ReLU(),
            nn.Linear(256, 448)
        )

        self.fc_decoder_value = nn.Sequential(
            nn.Linear(448, 448),
            nn.ReLU()
        )

        self.fc_decoder_policy = nn.Sequential(
            nn.Linear(448, 448),
            nn.ReLU()
        )

        self.predictor_action = nn.Linear(448, act_size)
        self.predictor_int_v = nn.Linear(448, 1)
        self.predictor_ext_v = nn.Linear(448, 1)

        # todo暂时不添加初始化 对比初始化前后的变化

    
    def _get_conv_out(self, shape):
        o = self.conv_encoder(torch.zeros(1, *shape))
        return int(np.prod(o.size()))
    

    def forward(self, x):
        x = x.float() / 255
        conv_out = self.conv_encoder(x)
        conv_out = conv_out.view(conv_out.size()[0], -1)
        fc_out = self.fc_encoder(conv_out)

        x_v = fc_out + self.fc_decoder_value(fc_out)
        x_pi = fc_out + self.fc_decoder_policy(fc_out)
        int_value = self.predictor_int_v(x_v).squeeze(-1)
        ext_value = self.predictor_ext_v(x_v).squeeze(-1)
        logits = self.predictor_action(x_pi)
        probs = F.softmax(logits, dim=1)
        dist = Categorical(probs)

        return dist, int_value, ext_value, probs
    

class PredictorModel(nn.Module):

    def __init__(self, obs_size):
        super().__init__()
        
        self.conv_encoder = nn.Sequential(
            nn.Conv2d(obs_size[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )


        self.conv_out_size = self._get_conv_out(obs_size)
        self.fc_decoder = nn.Sequential(
            nn.Linear(self.conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 512)
        )

    def _get_conv_out(self, shape):
        o = self.conv_encoder(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        conv_out = self.conv_encoder(x)
        fc_out =self.fc_decoder(conv_out.view(conv_out.size()[0], -1))
        return fc_out
    

class TargetModel(nn.Module):

    def __init__(self, obs_size):
        super().__init__()
        self.obs_size = obs_size

        self.conv_encoder = nn.Sequential(
            nn.Conv2d(obs_size[0], 32, kernel_size=8, stride=4),
            nn.LeakyReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.LeakyReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.LeakyReLU(),
        )

        self.conv_out_size = self._get_conv_out(obs_size)
        self.fc_decoder = nn.Sequential(
            nn.Linear(self.conv_out_size, 512)
        )

    def _get_conv_out(self, shape):
        o = self.conv_encoder(torch.zeros(1, *shape))
        return int(np.prod(o.size()))
    
    
    def forward(self, x):
        conv_out = self.conv_encoder(x)
        fc_out = self.fc_decoder(conv_out.view(conv_out.size()[0], -1))
        return fc_out